import numpy as np
from common.geometry.geom import Geom, bbox_to_geom

import geojson
import os
import datetime
import random

from catalog.place.place_api import PlaceAPI
from catalog.raster.raster import retrieve_by_query

from typing import Union


class EvalSegModule(object):
    def __init__(self, country_name="CHN") -> None:
        """Initialize an Evaluation Module of Image Segmentation.
        The initialization is restricted within the given country.

        Parameters
        ----------
        country_name : str, optional
            The short name of a given country, by default "CHN"

        Raises
        ------
        TypeError
            The data of CHN and USA are stored only, so no other country is accepted.
        """
        country_name_upper = country_name.upper()
        if country_name_upper == "CHN" or country_name_upper == "CHINA":
            country_fullname = "people's republic of china"
        elif country_name_upper == "USA" or country_name_upper == "US":
            country_fullname = "united states"
        else:
            raise TypeError("Only accepting CHN or USA right now...")

        country_aoi = PlaceAPI.get_geom_list_with_place_name(
            name_en=country_fullname, region_level=0
        )

        self.country_name = country_name_upper
        self.region_contour = country_aoi[0]

    def generate_random_aois(
        self,
        num_location: int,
        start_date: Union[int, str],
        end_date: Union[int, str],
        buffer_size=0.2,
        land_cover_type=5,
        land_cover_perc=0.5,
        save_file=False,
        output_folder=None,
    ):
        # TODO, this is a temp fix, to avoid running and importing this function in remote github CI/CD environment
        # otherwise that import would cause to load a gpkg file only lives in our NAS
        # we should later move away from `stats`
        from legacy.stats.utils_iolulc import retrieve_by_aoi

        region_contour_bounds = self.region_contour.bounds
        xmin = region_contour_bounds.left
        xmax = region_contour_bounds.right
        ymax = region_contour_bounds.upper
        ymin = region_contour_bounds.bottom

        temporal_period = (
            datetime.datetime.strptime(str(start_date), "%Y%m%d"),
            datetime.datetime.strptime(str(end_date), "%Y%m%d"),
        )

        output_path_list = []
        features = []
        counter = 0

        while counter < num_location:

            rand_x = random.uniform(xmin, xmax)
            rand_y = random.uniform(ymin, ymax)

            point_geom = Geom.make_point((rand_x, rand_y))
            point_geom_buffered = point_geom.buffer(buffer_size)
            bbox = point_geom_buffered.bounds
            bbox_geom = bbox_to_geom(bbox)

            if bbox_geom.within(self.region_contour):
                lulc_rds = retrieve_by_aoi(bbox_geom)
                lulc_arr = np.squeeze(lulc_rds.data)

                cnt = np.count_nonzero(lulc_arr == land_cover_type)  # crop
                lulc_perc = cnt / lulc_arr.shape[0] / lulc_arr.shape[1]

                if lulc_perc > land_cover_perc:
                    counter += 1

                    patch_rds = retrieve_by_query(
                        collection_id="COPERNICUS/SENTINEL2/L2A",
                        spatial=bbox_geom,
                        temporal=temporal_period,
                        asset_names=["B04", "B03", "B02", "B08"],
                        composite_strategy="s2decloud-v1",
                        dest_srs="EPSG:4326",
                    )

                    if save_file and output_folder is not None:
                        savepath = os.path.join(
                            output_folder, f"patch_for_seg_test_{counter}.tif"
                        )
                        patch_rds.to_geotiff(savepath)

                        bbox_contour = [
                            (bbox.left, bbox.upper),
                            (bbox.left, bbox.bottom),
                            (bbox.right, bbox.bottom),
                            (bbox.right, bbox.upper),
                        ]

                        bbox_corrected = tuple([bbox_contour])
                        points = geojson.Polygon(bbox_corrected)
                        features.append(geojson.Feature(id=counter, geometry=points))
                        output_path_list.append(savepath)

        if save_file and output_folder is not None:
            feature_collection = geojson.FeatureCollection(features)
            output_filepath = os.path.join(output_folder, "patch_for_seg_test.geojson")
            with open(output_filepath, "w") as f:
                geojson.dump(feature_collection, f, indent=2)

        return output_path_list


if __name__ == "__main__":
    output_folder = "/NAS6/tmp_zhicheng/data_2/TMP"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    tmp_cls = EvalSegModule()
    tmp_cls.generate_random_aois(
        num_location=3,
        start_date=20191001,
        end_date=20191101,
        save_file=True,
        output_folder=output_folder,
    )

    print("Done!")
